#include <stdio.h>
#include "curand_kernel.h"
#include "cublas_v2.h"

// Pointers to device memory
float *d_pdf_param1_current;
float *d_pdf_param2_current;
float *d_pdf_param1_new;
float *d_pdf_param2_new;
float *d_dx;
float *d_x;
float *d_p0_x;
float *d_p0_pdf_param1;
float *d_p0_pdf_param2;
float *d_p1_x;
float *d_p1_pdf_param1;
float *d_p1_pdf_param2;
int *d_reject;
int *d_reject_vector;
float *d_n0;
float *d_n0T;
float *d_n1T;
float *d_stepsize_pdf_param1;
float *d_stepsize_pdf_param2;
float *d_pdf_param1;
float *d_pdf_param2;
int *d_pdf_param1_pdf;
int *d_pdf_param2_pdf;
float *d_Tinv_current;
float *d_Sinv_current;
float *d_H;
float *d_C;
float *d_dy;
float *d_n1;
float *d_dum3;
int *d_x_pdf;
float *d_stepsize;
float *d_pdf_hyperparam1;
float *d_pdf_hyperparam2;
bool *d_accept;
float *d_one;
float *d_zero;

curandState *d_rngStates;
cublasHandle_t handle;

// Problem dimensions
int nmeasuremax, numsites, nmeasure, statesize, dim1, dim2, nmeasuretotal;

extern "C" void allocate_gpu_mem_cuda(int nmeasuremax_in, int numsites_in, int nmeasure_in,
                        int statesize_in, int dim1_in, int dim2_in, int nmeasuretotal_in)
{
   //printf("Allocating device memory...\n");

   nmeasuremax = nmeasuremax_in;
   numsites = numsites_in;
   nmeasure = nmeasure_in;
   statesize = statesize_in;
   dim1 = dim1_in;
   dim2 = dim2_in;
   nmeasuretotal = nmeasuretotal_in;

   int bytes = 0;
   int i = sizeof(int);
   int r = sizeof(float);
   int b = sizeof(bool);
   cudaError_t err;

   err = cudaMalloc (&d_pdf_param1_current, statesize*r);
   if (err){
     printf("ERROR: could not allocate pdf_param1_current\n");
     printf("Error was %d\n",err);
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_param1_new, statesize*r);
   if (err){
     printf("ERROR: could not allocate pdf_param1_new\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_param2_current, statesize*r);
   if (err){
     printf("ERROR: could not allocate pdf_param2_current\n");
     printf("Error was %d\n",err);
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_param2_new, statesize*r);
   if (err){
     printf("ERROR: could not allocate pdf_param2_new\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_dx, statesize*r);
   if (err){
     printf("ERROR: could not allocate dx\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_x, statesize*r);
   if (err){
     printf("ERROR: could not allocate x\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_p0_x, statesize*r);
   if (err){
     printf("ERROR: could not allocate p0_x\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_p0_pdf_param1, statesize*r);
   if (err){
     printf("ERROR: could not allocate p0_pdf_param1\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_p0_pdf_param2, statesize*r);
   if (err){
     printf("ERROR: could not allocate p0_pdf_param2\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_p1_x, statesize*r);
   if (err){
     printf("ERROR: could not allocate p1_x\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_p1_pdf_param1, statesize*r);
   if (err){
     printf("ERROR: could not allocate p1_pdf_param1\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_p1_pdf_param2, statesize*r);
   if (err){
     printf("ERROR: could not allocate p1_pdf_param2\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_reject, i);
   if (err){
     printf("ERROR: could not allocate reject\n");
     exit(1);
   }
   bytes = bytes + i;

   err = cudaMalloc (&d_reject_vector, statesize*i);
   if (err){
     printf("ERROR: could not allocate reject_vector\n");
     exit(1);
   }
   bytes = bytes + statesize * i;

   err = cudaMalloc (&d_n0, nmeasuremax * r);
   if (err){
     printf("ERROR: could not allocate n0\n");
     exit(1);
   }
   bytes = bytes + nmeasuremax * r;

   err = cudaMalloc (&d_n1, nmeasuremax * r);
   if (err){
     printf("ERROR: could not allocate n1\n");
     exit(1);
   }
   bytes = bytes + nmeasuremax * r;

   err = cudaMalloc (&d_n0T, r);
   if (err){
     printf("ERROR: could not allocate n0T\n");
     exit(1);
   }
   bytes = bytes + r;

   err = cudaMalloc (&d_n1T, r);
   if (err){
     printf("ERROR: could not allocate n1T\n");
     exit(1);
   }
   bytes = bytes + r;

   err = cudaMalloc (&d_stepsize_pdf_param1, statesize * r);
   if (err){
     printf("ERROR: could not allocate stepsize_pdf_param1\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_stepsize_pdf_param2, statesize * r);
   if (err){
     printf("ERROR: could not allocate stepsize_pdf_param2\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_param1, statesize * r);
   if (err){
     printf("ERROR: could not allocate pdf_param1\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_param2, statesize * r);
   if (err){
     printf("ERROR: could not allocate pdf_param2\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_param1_pdf, statesize * i);
   if (err){
     printf("ERROR: could not allocate pdf_param1_pdf\n");
     exit(1);
   }
   bytes = bytes + statesize * i;

   err = cudaMalloc (&d_pdf_param2_pdf, statesize * i);
   if (err){
     printf("ERROR: could not allocate pdf_param2_pdf\n");
     exit(1);
   }
   bytes = bytes + statesize * i;

   err = cudaMalloc (&d_Tinv_current, nmeasuretotal * nmeasuretotal * r);
   if (err){
     printf("ERROR: could not allocate Tinv_current\n");
     exit(1);
   }
   bytes = bytes + nmeasuretotal * nmeasuretotal * r;

   err = cudaMalloc (&d_Sinv_current, numsites * numsites * r);
   if (err){
     printf("ERROR: could not allocate Sinv_current\n");
     exit(1);
   }
   bytes = bytes + numsites * numsites * r;

   err = cudaMalloc (&d_H, nmeasuremax * statesize * r);
   if (err){
     printf("ERROR: could not allocate H\n");
     exit(1);
   }
   bytes = bytes + nmeasuremax * statesize * r;

   err = cudaMalloc (&d_C, nmeasuremax * r);
   if (err){
     printf("ERROR: could not allocate C\n");
     exit(1);
   }
   bytes = bytes + nmeasuremax * r;

   err = cudaMalloc (&d_dy, nmeasuremax * r);
   if (err){
     printf("ERROR: could not allocate dy\n");
     exit(1);
   }
   bytes = bytes + nmeasuremax * r;

   err = cudaMalloc (&d_dum3, nmeasuretotal * numsites * r);
   if (err){
     printf("ERROR: could not allocate dum3\n");
     exit(1);
   }
   bytes = bytes + nmeasuretotal * numsites * r;

   err = cudaMalloc (&d_x_pdf, statesize*i);
   if (err){
     printf("ERROR: could not allocate x_pdf\n");
     exit(1);
   }
   bytes = bytes + statesize * i;

   err = cudaMalloc (&d_stepsize, statesize*r);
   if (err){
     printf("ERROR: could not allocate stepsize\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_hyperparam1, statesize*r);
   if (err){
     printf("ERROR: could not allocate pdf_hyperparam1\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_pdf_hyperparam2, statesize*r);
   if (err){
     printf("ERROR: could not allocate pdf_hyperparam2\n");
     exit(1);
   }
   bytes = bytes + statesize * r;

   err = cudaMalloc (&d_accept, b);
   if (err){
     printf("ERROR: could not allocate accept\n");
     exit(1);
   }
   bytes = bytes + b;

   err = cudaMalloc (&d_rngStates, nmeasuremax*sizeof(curandState));
   if (err){
     printf("ERROR: could not allocate rngStates\n");
     exit(1);
   }
   bytes = bytes + nmeasuremax*sizeof(curandState);

   err = cudaMalloc (&d_one, r);
   if (err){
     printf("ERROR: could not allocate one\n");
     exit(1);
   }
   bytes = bytes + r;

   err = cudaMalloc (&d_zero, r);
   if (err){
     printf("ERROR: could not allocate zero\n");
     exit(1);
   }
   bytes = bytes + r;

   printf(" Allocated : %d MB on GPU\n",bytes/1024/1024);
}

// RNG init kernel
__global__ void initRNG(curandState *rngStates, const int nmeasuremax,
                        const unsigned int seed)
{
   // Determine thread ID
   unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;

   // Initialise the RNG
   if (tid < nmeasuremax)
   {
      curand_init(seed, tid, 0, &rngStates[tid]);
   }
}

__global__ void init_constants(float *one, float *zero)
{
   *one = 1.0f;
   *zero = 0.0f;
}

extern "C" void setup_gpu_cuda(float *pdf_param1_current, float *pdf_param2_current, float *x,
                    float *p0_x, float *p0_pdf_param1, float *p0_pdf_param2,
                    float *stepsize_pdf_param1, float *stepsize_pdf_param2,
                    float *pdf_param1, float *pdf_param2, int *pdf_param1_pdf,
                    int *pdf_param2_pdf, float *stepsize, float *pdf_hyperparam1,
                    float *pdf_hyperparam2, float *H, int *x_pdf)
{
   //printf("Setting up device with initial data...\n");

   int i = sizeof(int);
   int r = sizeof(float);
   cudaError_t err;

   err = cudaMemcpy(d_pdf_param1_current, pdf_param1_current, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_param1_current to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_pdf_param2_current, pdf_param2_current, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_param2_current to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_x, x, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy x to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_p0_x, p0_x, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy p0_x to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_p0_pdf_param1, p0_pdf_param1, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy p0_pdf_param1 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_p0_pdf_param2, p0_pdf_param2, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy p0_pdf_param2 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_stepsize_pdf_param1, stepsize_pdf_param1, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy stepsize_pdf_param1 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_stepsize_pdf_param2, stepsize_pdf_param2, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy stepsize_pdf_param2 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_pdf_param1, pdf_param1, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_param1 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_pdf_param2, pdf_param2, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_param2 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_pdf_param1_pdf, pdf_param1_pdf, statesize*i, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_param1_pdf to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_pdf_param2_pdf, pdf_param2_pdf, statesize*i, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_param2_pdf to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_stepsize, stepsize, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy stepsize to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_pdf_hyperparam1, pdf_hyperparam1, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_hyperparam1 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_pdf_hyperparam2, pdf_hyperparam2, statesize*r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy pdf_hyperparam2 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_H, H, nmeasuremax*statesize*i, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy H to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_x_pdf, x_pdf, statesize*i, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy x_pdf to device\n");
     exit(1);
   }

   // Zero reject counts

   err = cudaMemset(d_reject, 0, i);
   if (err){
     printf("ERROR: could not zero reject on device\n");
     exit(1);
   }

   err = cudaMemset(d_reject_vector, 0, statesize*i);
   if (err){
     printf("ERROR: could not zero reject_vector on device\n");
     exit(1);
   }

   //printf("Done\n");

   // Fixed mapping of threads to blocks

   dim3 block, grid; 
   int threads_per_block = 128;
   int seed = 1234;
   block.x = threads_per_block;
   grid.x = (nmeasuremax + threads_per_block - 1) / threads_per_block;
   initRNG<<<grid, block>>>(d_rngStates, nmeasuremax, seed);

   //printf("Initialized GPU RNG\n");

   init_constants<<<1,1>>>(d_one, d_zero);

   cublasCreate(&handle);
   cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_DEVICE);
}

extern "C" void copy_data_back_cuda(float *pdf_param1_current, float *pdf_param2_current,
                         float *x, int *reject)
{
   //printf("Copying data back from device\n");

   int err;
   int i = sizeof(int);
   int r = sizeof(float);

   err = cudaMemcpy(pdf_param1_current, d_pdf_param1_current, statesize * r, cudaMemcpyDeviceToHost);
   if (err){
     printf("ERROR: could not copy pdf_param1_current from device\n");
     exit(1);
   }

   err = cudaMemcpy(pdf_param2_current, d_pdf_param2_current, statesize * r, cudaMemcpyDeviceToHost);
   if (err){
     printf("ERROR: could not copy pdf_param2_current from device\n");
     exit(1);
   }

   err = cudaMemcpy(x, d_x, statesize * r, cudaMemcpyDeviceToHost);
   if (err){
     printf("ERROR: could not copy x from device\n");
     exit(1);
   }

   err = cudaMemcpy(reject, d_reject, i, cudaMemcpyDeviceToHost);
   if (err){
     printf("ERROR: could not copy reject from device\n");
     exit(1);
   }
}

extern "C" void copy_final_data_back_cuda(int *reject_vector)
{
   //printf("Copying final data back from device\n");

   int err;
   int i = sizeof(int);

   err = cudaMemcpy(reject_vector, d_reject_vector, statesize*i, cudaMemcpyDeviceToHost);
   if (err){
     printf("ERROR: could not copy reject_vector from device\n");
     exit(1);
   }
}

extern "C" void free_gpu_mem_cuda()
{
   //printf("Freeing device memory...\n");

   cudaFree(d_pdf_param1_current);
   cudaFree(d_pdf_param2_current);
   cudaFree(d_pdf_param1_new);
   cudaFree(d_pdf_param2_new);
   cudaFree(d_dx);
   cudaFree(d_x);
   cudaFree(d_p0_x);
   cudaFree(d_p0_pdf_param1);
   cudaFree(d_p0_pdf_param2);
   cudaFree(d_p1_x);
   cudaFree(d_p1_pdf_param1);
   cudaFree(d_p1_pdf_param2);
   cudaFree(d_reject);
   cudaFree(d_reject_vector);
   cudaFree(d_n0);
   cudaFree(d_n1);
   cudaFree(d_n0T);
   cudaFree(d_n1T);
   cudaFree(d_stepsize_pdf_param1);
   cudaFree(d_stepsize_pdf_param2);
   cudaFree(d_pdf_param1);
   cudaFree(d_pdf_param2);
   cudaFree(d_pdf_param1_pdf);
   cudaFree(d_pdf_param2_pdf);
   cudaFree(d_Tinv_current);
   cudaFree(d_Sinv_current);
   cudaFree(d_H);
   cudaFree(d_C);
   cudaFree(d_dy);
   cudaFree(d_dum3);
   cudaFree(d_stepsize);
   cudaFree(d_x_pdf);
   cudaFree(d_pdf_hyperparam1);
   cudaFree(d_pdf_hyperparam2);
   cudaFree(d_accept);
   cudaFree(d_rngStates);
   cudaFree(d_one);
   cudaFree(d_zero);

   cublasDestroy(handle);

   //printf("Done\n");
}

#define SQRT_2_PI 2.506628274794649f

__device__ float pdf_calc(float x, float pdf_param1, float pdf_param2, int pdf)
{
   float p1,tmp;
   bool positive;

   // PDF is LOGNORMAL
   if (pdf == 0)
   {
      tmp = logf(x) - logf(pdf_param1);
      p1 = (1.0f / (x * pdf_param2 * SQRT_2_PI)) * expf( -1.0f * tmp * tmp / 2.0f / (pdf_param2 * pdf_param2));
      positive = true;
   }

   // PDF is EXPONENTIAL
   if (pdf == 1)
   {
      p1 = pdf_param1 * expf(-1.0f * pdf_param1 * x);
      positive = true;
   }

   // PDF is GAUSSIAN
   if (pdf == 2)
   {
      tmp = x - pdf_param1;
      p1 = (1.0f / (pdf_param2 * SQRT_2_PI)) * expf( -1.0f * tmp * tmp / 2.0 / (pdf_param2 * pdf_param2));
      positive = false;
   }

   // PDF is UNIFORM
   if (pdf == 3)
   {
      if ( x < pdf_param1 || x > pdf_param2)
      {
         p1 = 0.0f;
      }
      else
      {
         p1 = 1.0f / (pdf_param2 - pdf_param1);
      }
      positive = false;
   }

   if ( x <= 0.0f && positive ) p1 = 0.0f;

   return p1;
}

__global__ void compute_x_delta(curandState *rngStates, const int statesize,
        float *stepsize_pdf_param1, float *pdf_param1, int *pdf_param1_pdf,
        float *pdf_param1_current, float *pdf_param1_new, float *pdf_hyperparam1, float *stepsize_pdf_param2, 
        float *pdf_param2, int *pdf_param2_pdf, float *pdf_param2_current,
        float *pdf_param2_new, float *pdf_hyperparam2, float *stepsize, float *dx,
        float *p1_pdf_param1, float *p1_pdf_param2, float *p1_x, float *x, int *x_pdf)
{
   // Determine thread ID
   unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
   if (tid >= statesize) return;

   // Compute the new values via RNG
   curandState state = rngStates[tid];
   pdf_param1_new[tid] = pdf_param1_current[tid] + stepsize_pdf_param1[tid] * curand_normal(&state);
   pdf_param2_new[tid] = pdf_param2_current[tid] + stepsize_pdf_param2[tid] * curand_normal(&state);
   dx[tid] = stepsize[tid] * curand_normal(&state);
   rngStates[tid] = state;

   // Compute the corresponding probabilities
   p1_pdf_param1[tid] = pdf_calc(pdf_param1_new[tid], pdf_param1[tid], pdf_hyperparam1[tid], pdf_param1_pdf[tid]);
   p1_pdf_param2[tid] = pdf_calc(pdf_param2_new[tid], pdf_param2[tid], pdf_hyperparam2[tid], pdf_param2_pdf[tid]);
   p1_x[tid] = pdf_calc(x[tid]+dx[tid], pdf_param1_new[tid], pdf_param2_new[tid], x_pdf[tid]);
}

__global__ void compute_x_n1(const int nmeasuremax,
                             float *n1, float *n0, float *H, float *dx)
{
   // Determine thread ID
   unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
   if (tid >= nmeasuremax) return;

   n1[tid] = n0[tid] + H[tid] * (*dx);
}

__global__ void compute_x_accept(curandState *rngStates, float *x, float *dx,
                         float *n0T, float *n1T,
                         float *p0_x, float *p0_pdf_param1, float *p0_pdf_param2,
                         float *p1_x, float *p1_pdf_param1, float *p1_pdf_param2,
                         float *pdf_param1_current, float *pdf_param1_new,
                         float *pdf_param2_current, float *pdf_param2_new,
                         bool *accept, bool burn_in_done, int *reject, int *reject_vector)
{
   // Compute P1/P0
   float pT = logf ( ( *p1_x * *p1_pdf_param1 * *p1_pdf_param2 ) / 
                     ( *p0_x * *p0_pdf_param1 * *p0_pdf_param2 ) ) -
                      0.5f * ( *n1T - *n0T);

   // Generate uniform random

   curandState state = rngStates[0];
   float u = curand_uniform(&state);
   rngStates[0] = state;

   // Compare to acceptance probability

   *accept = (logf(u) <= pT);

   if (*accept)
   {
      // Accept
      *p0_x = *p1_x;
      *p0_pdf_param1 = *p1_pdf_param1;
      *p0_pdf_param2 = *p1_pdf_param2;
      *x = *x + *dx;
      *pdf_param1_current = *pdf_param1_new;
      *pdf_param2_current = *pdf_param2_new;
      *n0T = *n1T;
   }
   else
   {
      // Reject
      if (burn_in_done)
      {
         *reject = *reject + 1;
         *reject_vector = *reject_vector + 1;
      }
   }

}

__global__ void compute_x_array_update(const int nmeasuremax, bool *accept,
                                     float *n0, float *n1)
{
   // Determine thread ID 
   unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
   if (tid >= nmeasuremax) return;

   if (*accept)
   {
      n0[tid] = n1[tid]; 
   }
}

extern "C" void compute_x_gpu_cuda(float *n0T, float *n0,
                                   float *Tinv_current, float *Sinv_current,
                                   bool burn_in_done)
{
   //printf("Copy input data onto GPU\n");

   int r = sizeof(float);
   cudaError_t err;

   err = cudaMemcpy(d_n0T, n0T, r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy n0T to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_n0, n0, nmeasuremax * r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy n0 to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_Tinv_current, Tinv_current, nmeasuretotal * nmeasuretotal * r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy Tinv_current to device\n");
     exit(1);
   }

   err = cudaMemcpy(d_Sinv_current, Sinv_current, numsites * numsites * r, cudaMemcpyHostToDevice);
   if (err){
     printf("ERROR: could not copy Sinv_current to device\n");
     exit(1);
   }
  
   //printf("Compute_x on GPU\n");

   // Fixed mapping of threads to blocks

   dim3 block, grid;
   int threads_per_block = 128;
   block.x = threads_per_block;

   // First kernel : compute the proposed deltas

   grid.x = (statesize + threads_per_block - 1) / threads_per_block;

   compute_x_delta<<<grid, block>>>(d_rngStates,statesize,
        d_stepsize_pdf_param1,d_pdf_param1,
        d_pdf_param1_pdf,d_pdf_param1_current,d_pdf_param1_new,d_pdf_hyperparam1,
        d_stepsize_pdf_param2,d_pdf_param2,d_pdf_param2_pdf,
        d_pdf_param2_current, d_pdf_param2_new,d_pdf_hyperparam2,d_stepsize,d_dx,
        d_p1_pdf_param1, d_p1_pdf_param2, d_p1_x, d_x, d_x_pdf);

   // Now work through each element in sequence

   grid.x = (nmeasuremax + threads_per_block - 1) / threads_per_block;

   for (int xi = 0; xi < statesize; xi++)
   {
  
      // Second kernel : compute the change n0->n1
      // Pass in the xi'th column of H (stored contiguous, F90 col-major ordering)
      compute_x_n1<<<grid, block>>>(nmeasuremax, d_n1, d_n0, &d_H[xi*nmeasuremax], &d_dx[xi]);

      // Now compute the products
      // dum3 = Tinv_current * n1
      cublasSgemm(handle,CUBLAS_OP_N,CUBLAS_OP_N, nmeasuretotal, numsites, nmeasuretotal, d_one, d_Tinv_current, nmeasuretotal, d_n1, nmeasuretotal, d_zero, d_dum3, nmeasuretotal);

      // C = dum3 * Sinv_current
      cublasSgemm(handle,CUBLAS_OP_N,CUBLAS_OP_N, nmeasuretotal, numsites, numsites, d_one, d_dum3, nmeasuretotal, d_Sinv_current, numsites, d_zero, d_C, nmeasuretotal);

      // n1T = n1 . C
      cublasSdot(handle, nmeasuremax, d_n1, 1, d_C, 1, d_n1T); 

      // Third kernel : Compute the update probability and randomly accept/reject
      compute_x_accept<<<1,1>>>(d_rngStates, &d_x[xi], &d_dx[xi],
                    d_n0T, d_n1T,
                    &d_p0_x[xi], &d_p0_pdf_param1[xi], &d_p0_pdf_param2[xi],
                    &d_p1_x[xi], &d_p1_pdf_param1[xi], &d_p1_pdf_param2[xi],
                    &d_pdf_param1_current[xi], &d_pdf_param1_new[xi],
                    &d_pdf_param2_current[xi], &d_pdf_param2_new[xi],
                    d_accept, burn_in_done , d_reject, &d_reject_vector[xi] );

      // Fourth kernel : Update n0 arrays if needed
      compute_x_array_update<<<grid, block>>>(nmeasuremax, d_accept,
                                     d_n0, d_n1);

   }

   //
   //printf("Copy outputs back from GPU\n");

   err = cudaMemcpy(n0T, d_n0T, r, cudaMemcpyDeviceToHost);
   if (err){
     printf("ERROR: could not copy n0T from device\n");
     exit(1);
   }

   err = cudaMemcpy(n0, d_n0, nmeasuremax * r, cudaMemcpyDeviceToHost);
   if (err){
     printf("ERROR: could not copy n0 from device\n");
     exit(1);
   }
}
